#include <ntddk.h>
#include <wdm.h>

#define DEVICE_NAME L"\\Device\\DeviceA"
#define SYMBOLICLINK_NAME L"\\DosDevices\\DeviceA"

#ifdef DBG
#define ENTER() DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Debug] Callee: Enter %s\r\n", __FUNCTION__)
#define LEAVE() DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Debug] Callee: Leave %s\r\n", __FUNCTION__)
#define LOG(format, ...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Debug] %s " format, __FUNCTION__, __VA_ARGS__)
#define ERR(format, ...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_ERROR_LEVEL, "[Error] %s " format, __FUNCTION__, __VA_ARGS__)
#else
#define ENTER()
#define LEAVE()
#define LOG(format, ...)
#define ERR(format, ...)
#endif // DBG

typedef struct _DEVICE_EXTENSION
{
    PDEVICE_OBJECT DeviceObject;
    UNICODE_STRING DeviceName;
    UNICODE_STRING SymbolicLinkName;
    KTIMER PollingTimer;
    KDPC PollingDpc;
    PIRP PengdingIrp;
}DEVICE_EXTENSION, *PDEVICE_EXTENSION;

NTSTATUS
DriverCreate(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
);

NTSTATUS
DriverClose(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
);

NTSTATUS
DriverRead(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
);

VOID
DriverUnload(
    _In_ PDRIVER_OBJECT DriverObject
);

NTSTATUS
DeviceCreate(
    _In_ PDRIVER_OBJECT DriverObject
);

VOID
OnTimerDpc(
    _In_ struct _KDPC* Dpc,
    _In_opt_ PVOID DeferredContext,
    _In_opt_ PVOID SystemArgument1,
    _In_opt_ PVOID SystemArgument2
);

NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT DriverObject,
    _In_ PUNICODE_STRING RegistryPath
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(RegistryPath);
    
    ENTER(); 
    DriverObject->MajorFunction[IRP_MJ_CREATE] = DriverCreate;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverClose;
    DriverObject->MajorFunction[IRP_MJ_READ] = DriverRead;
    DriverObject->DriverUnload = DriverUnload;
    status = DeviceCreate(DriverObject);
    LEAVE();

    return status;
}

NTSTATUS
DriverCreate(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
    )
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);
    ENTER();
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LEAVE();
    return status;
}

NTSTATUS
DriverClose(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
)
{
    NTSTATUS status;

    UNREFERENCED_PARAMETER(DeviceObject);
    ENTER();
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LEAVE();
    return status;
}

NTSTATUS
DriverRead(
    _In_ PDEVICE_OBJECT DeviceObject,
    _In_ PIRP Irp
)
{
    PDEVICE_EXTENSION pDeviceExtension;
    LARGE_INTEGER dueTime;
    
    ENTER();
    pDeviceExtension = DeviceObject->DeviceExtension;
    if (!pDeviceExtension->PengdingIrp)
    {
        pDeviceExtension->PengdingIrp = Irp;
    }
    IoMarkIrpPending(Irp);
    dueTime.QuadPart = 3 * -1 * 10000000;
    KeSetTimer(&pDeviceExtension->PollingTimer, dueTime, &pDeviceExtension->PollingDpc);
    LEAVE();

    return STATUS_PENDING;
}

VOID
DriverUnload(
    _In_ PDRIVER_OBJECT DriverObject
)
{
    NTSTATUS status;
    PDEVICE_OBJECT pDeviceObject, pNextDeviceObject;
    PDEVICE_EXTENSION pDeviceExtension;
    
    ENTER();
    pDeviceObject = DriverObject->DeviceObject;
    while (pDeviceObject)
    {
        pNextDeviceObject = pDeviceObject->NextDevice;
        pDeviceExtension = pDeviceObject->DeviceExtension;

        LOG("IoDeleteSymbolicLink DeviceObject=%p, DeviceName:%wZ, SymbolicLinkName=%wZ\r\n", 
            pDeviceExtension->DeviceObject, 
            &pDeviceExtension->DeviceName, 
            &pDeviceExtension->SymbolicLinkName);
        status = IoDeleteSymbolicLink(&pDeviceExtension->SymbolicLinkName);
        if (!NT_SUCCESS(status))
        {
            ERR("IoDeleteSymbolicLink Symbolic Link Name:%wZ, status=%#X\r\n", &pDeviceExtension->SymbolicLinkName, status);
        }
        IoDeleteDevice(pDeviceObject);
        pDeviceObject = pNextDeviceObject;
    }
    LEAVE();
}

NTSTATUS
DeviceCreate(
    _In_ PDRIVER_OBJECT DriverObject
)
{
    NTSTATUS status;
    UNICODE_STRING deviceName;
    UNICODE_STRING symbolicLinkName;
    PDEVICE_OBJECT pDeviceObject;
    PDEVICE_EXTENSION pDeviceExtension;

    ENTER();
    __try
    {
        RtlInitUnicodeString(&deviceName, DEVICE_NAME);
        status = IoCreateDevice(
            DriverObject,
            sizeof(DEVICE_EXTENSION),
            &deviceName,
            FILE_DEVICE_UNKNOWN,
            FILE_DEVICE_SECURE_OPEN,
            FALSE,
            &pDeviceObject
        );
        if (!NT_SUCCESS(status))
        {
            ERR("IoCreateDevice status=%#X\r\n", status);
            __leave;
        }
        pDeviceObject->Flags |= DO_BUFFERED_IO;

        RtlInitUnicodeString(&symbolicLinkName, SYMBOLICLINK_NAME);
        status = IoCreateSymbolicLink(&symbolicLinkName, &deviceName);
        if (!NT_SUCCESS(status))
        {
            ERR("IoCreateSymbolicLink status=%#X\r\n", status);
            IoDeleteDevice(pDeviceObject);
            __leave;
        }

        pDeviceExtension = pDeviceObject->DeviceExtension;
        pDeviceExtension->DeviceObject = pDeviceObject;
        pDeviceExtension->DeviceName = deviceName;
        pDeviceExtension->SymbolicLinkName = symbolicLinkName;
        pDeviceExtension->PengdingIrp = NULL;
        KeInitializeTimer(&pDeviceExtension->PollingTimer);
        KeInitializeDpc(&pDeviceExtension->PollingDpc, OnTimerDpc, pDeviceObject);

        LOG("DeviceObject=%p, DeviceName:%wZ, SymbolicLinkName=%wZ\r\n",
            pDeviceExtension->DeviceObject,
            &pDeviceExtension->DeviceName,
            &pDeviceExtension->SymbolicLinkName);
    }
    __finally
    {
        LEAVE();
    }
    return status;
}

VOID
OnTimerDpc(
    _In_ struct _KDPC* Dpc,
    _In_opt_ PVOID DeferredContext,
    _In_opt_ PVOID SystemArgument1,
    _In_opt_ PVOID SystemArgument2
)
{
    PDEVICE_OBJECT pDeviceObject;
    PDEVICE_EXTENSION pDeviceExtension;

    UNREFERENCED_PARAMETER(Dpc);
    UNREFERENCED_PARAMETER(SystemArgument1);
    UNREFERENCED_PARAMETER(SystemArgument2);
    
    ENTER();
    pDeviceObject = (PDEVICE_OBJECT)DeferredContext;
    pDeviceExtension = pDeviceObject->DeviceExtension;
    if (pDeviceExtension->PengdingIrp)
    {
        pDeviceExtension->PengdingIrp->IoStatus.Status = STATUS_SUCCESS;
        pDeviceExtension->PengdingIrp->IoStatus.Information = 0;
        IoCompleteRequest(pDeviceExtension->PengdingIrp, IO_NO_INCREMENT);
        pDeviceExtension->PengdingIrp = NULL;
    }
    LEAVE();
}